use crate::page::Page;
use crate::size_class::{PageKind, SizeClass};
use libc::{self, c_void, MAP_ANONYMOUS, MAP_PRIVATE, PROT_READ, PROT_WRITE};

pub const SEGMENT_SIZE: usize = 16 * 1024 * 1024; // 16MB per segment

#[derive(Debug)]
pub struct Segment {
    start: *mut u8,           // start of segment memory
    pages: Vec<Option<Page>>, // array of pages
    used: usize,              // number of pages in use
    capacity: usize,          // total number of pages
    page_kind: PageKind,      // type of pages in segment
    page_shift: usize,        // page size = 1 << page_shift
    segment_info_size: usize, // size of segment metadata
    prev: *mut Segment,       // previous segment in queue
    next: *mut Segment,       // next segment in queue
}

impl Segment {
    pub fn new(thread_id: usize, page_kind: PageKind) -> Self {
        debug_log!("\n🧩 Creating new segment:");
        debug_log!("  ├─ Thread ID: {}", thread_id);
        debug_log!("  ├─ Page kind: {:?}", page_kind);

        // Calculate page shift based on page kind
        let page_shift = match page_kind {
            PageKind::Small => 16,  // 64KB
            PageKind::Medium => 19, // 512KB
            PageKind::Large => 22,  // 4MB
        };

        // Calculate number of pages that can fit in a segment
        let capacity = match page_kind {
            PageKind::Small => 256, // 16MB/64KB = 256 pages
            PageKind::Medium => 32, // 16MB/512KB = 32 pages
            PageKind::Large => 4,   // 16MB/4MB = 4 pages
        };

        // Allocate segment memory using mmap
        let start = unsafe {
            let ptr = libc::mmap(
                std::ptr::null_mut(),
                SEGMENT_SIZE,
                PROT_READ | PROT_WRITE,
                MAP_PRIVATE | MAP_ANONYMOUS,
                -1,
                0,
            );

            if ptr == libc::MAP_FAILED {
                panic!("mmap failed");
            }

            ptr as *mut u8
        };

        debug_log!("  ├─ Page size: {} KB", 1 << (page_shift - 10));
        debug_log!("  ├─ Capacity: {} pages", capacity);
        debug_log!(
            "  ├─ Metadata size: {} bytes",
            std::mem::size_of::<Segment>()
        );
        debug_log!("  └─ Start address: {:?}", start);

        // Initialize segment metadata
        let segment_info_size =
            std::mem::size_of::<Segment>() + (capacity * std::mem::size_of::<Page>());
        let aligned_info_size = (segment_info_size + 15) & !15; // align to 16 bytes

        let mut pages = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            pages.push(None);
        }

        Segment {
            start,
            pages,
            used: 0,
            capacity,
            page_kind,
            page_shift,
            segment_info_size: aligned_info_size,
            prev: std::ptr::null_mut(),
            next: std::ptr::null_mut(),
        }
    }

    // Get the start address of a page area, accounting for metadata in first page
    pub fn get_page_area_start(&self, index: usize) -> *mut u8 {
        if index >= self.capacity {
            return std::ptr::null_mut();
        }

        unsafe {
            if index == 0 {
                match self.page_kind {
                    PageKind::Large => {
                        // For large pages, align metadata to 16 bytes
                        let aligned_info_size = (self.segment_info_size + 15) & !15;
                        self.start.add(aligned_info_size)
                    }
                    _ => self.start.add(self.segment_info_size),
                }
            } else {
                self.start.add(index * (1 << self.page_shift))
            }
        }
    }

    pub fn get_page_area_size(&self, index: usize) -> usize {
        if index >= self.capacity {
            return 0;
        }

        let page_size = 1 << self.page_shift;
        if index == 0 {
            match self.page_kind {
                PageKind::Large => {
                    // For large pages, use the entire segment minus aligned metadata
                    let _aligned_info_size = (self.segment_info_size + 15) & !15; // 16-byte align
                    debug_log!("  ├─ Aligned info size: {} bytes", _aligned_info_size);
                    SEGMENT_SIZE - _aligned_info_size
                }
                _ => {
                    // First page is smaller due to metadata
                    let _aligned_info_size = (self.segment_info_size + 15) & !15; // 16-byte align
                    debug_log!("  ├─ Aligned info size: {} bytes", _aligned_info_size);
                    debug_log!(
                        "  └─ Available area: {} bytes",
                        page_size - self.segment_info_size
                    );
                    page_size - self.segment_info_size
                }
            }
        } else {
            page_size
        }
    }

    // Initialize a page at the given index
    pub fn init_page(&mut self, index: usize, size_class: SizeClass) -> bool {
        if index >= self.capacity || self.pages[index].is_some() {
            return false;
        }

        let area_start = self.get_page_area_start(index);
        let area_size = self.get_page_area_size(index);

        // Create the page metadata
        let page = Page::new(area_start, area_size, size_class, index);
        self.pages[index] = Some(page);
        self.used += 1;
        true
    }

    pub fn has_free(&self) -> bool {
        self.used < self.capacity
    }

    pub fn find_free(&mut self) -> Option<usize> {
        for i in 0..self.capacity {
            if self.pages[i].is_none() {
                return Some(i);
            }
        }
        None
    }

    pub fn get_page(&mut self, index: usize) -> Option<&mut Page> {
        if index >= self.capacity {
            return None;
        }
        if let Some(ref mut page) = self.pages[index] {
            Some(page)
        } else {
            None
        }
    }

    pub fn is_page_initialized(&self, index: usize) -> bool {
        index < self.capacity && self.pages[index].is_some()
    }

    pub fn reclaim_page(&mut self, index: usize) -> bool {
        if index >= self.capacity {
            return false;
        }

        if let Some(page) = self.pages[index].take() {
            if !page.is_empty() {
                // Put it back if not empty
                self.pages[index] = Some(page);
                return false;
            }

            // Page is empty, can be reclaimed
            // No need to free page area as it's part of the segment memory
            self.used -= 1;
            return true;
        }
        false
    }

    pub fn try_reclaim_pages(&mut self) -> bool {
        let mut reclaimed = false;
        for i in 0..self.capacity {
            if self.reclaim_page(i) {
                reclaimed = true;
            }
        }
        reclaimed
    }

    pub fn is_empty(&self) -> bool {
        self.used == 0
    }

    pub fn get_next(&self) -> *mut Segment {
        self.next
    }

    pub fn get_prev(&self) -> *mut Segment {
        self.prev
    }

    pub fn set_next(&mut self, next: *mut Segment) {
        self.next = next;
    }

    pub fn set_prev(&mut self, prev: *mut Segment) {
        self.prev = prev;
    }
}

impl Drop for Segment {
    fn drop(&mut self) {
        unsafe {
            libc::munmap(self.start as *mut c_void, SEGMENT_SIZE);
        }
    }
}

impl std::fmt::Display for Segment {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Segment[{:?}] {{pages: {}/{}, kind: {}, addr: {:?}}}",
            self.start, self.used, self.capacity, self.page_kind, self.start
        )
    }
}
